{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run upon export from spreadsheet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total entries: 590\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "from astroquery.mast import Catalogs\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tces_file = '/mnt/tess/labels/s33_cam2ccd14_sample.csv'\n",
    "ext_data_file = '/mnt/tess/labels/ext_mast_data.csv'\n",
    "\n",
    "\n",
    "tce_table = pd.read_csv(tces_file, header=0, low_memory=False)\n",
    "tce_table['tic_id'] = tce_table['star_tic']\n",
    "tce_table['Duration'] = tce_table['planet_tdur']\n",
    "tce_table['Period'] = tce_table['planet_period']\n",
    "tce_table['RA'] = tce_table['star_ra']\n",
    "tce_table['Sectors'] = tce_table['sector_id'].apply(lambda v: len(v.split(' ')))\n",
    "tce_table['Transit_Depth'] = tce_table['planet_depth']\n",
    "tce_table['Dec'] = tce_table['star_dec']\n",
    "tce_table['teff'] = tce_table['star_teff']\n",
    "tce_table['SN'] = tce_table['snr']\n",
    "tce_table['Qingress'] = 0.0\n",
    "tce_table['Tmag'] = tce_table['star_tmag']\n",
    "tce_table['logg'] = tce_table['star_logg']\n",
    "tce_table['Epoc'] = tce_table['planet_epoch']\n",
    "tce_table = tce_table.set_index('tic_id')\n",
    "tce_table = tce_table.drop(columns=['Unnamed: 0'])\n",
    "\n",
    "tce_table['Duration'] /= 24.0\n",
    "\n",
    "# Drop some common invalid examples.\n",
    "# Orbits falling inside the star\n",
    "tce_table = tce_table[~tce_table.Ilabel]\n",
    "# Excessively large durations\n",
    "tce_table = tce_table[tce_table.Duration < 0.9 * tce_table.Period]\n",
    "\n",
    "joined_table = tce_table\n",
    "\n",
    "ext_table = pd.read_csv(ext_data_file, header=0, low_memory=False).set_index('tic_id')\n",
    "joined_table = joined_table.join(ext_table, on='tic_id', how='left')\n",
    "\n",
    "joined_table = joined_table[\n",
    "    joined_table['objType'].isnull()\n",
    "    | (joined_table['objType'] == 'STAR')\n",
    "]\n",
    "\n",
    "joined_table = joined_table.reset_index()[[\n",
    "    'tic_id', 'RA', 'Dec', 'Tmag', 'Epoc', 'Period', 'Duration',\n",
    "    'Transit_Depth', 'Sectors', 'star_rad', 'star_mass', 'teff',\n",
    "    'logg', 'SN', 'Qingress'\n",
    "]]\n",
    "\n",
    "\n",
    "disps = ['E', 'J', 'N', 'S', 'B']\n",
    "\n",
    "for d in disps:\n",
    "    joined_table[f'disp_{d}'] = 0\n",
    "\n",
    "joined_table = joined_table.set_index('tic_id')\n",
    "print(f'Total entries: {len(joined_table)}')\n",
    "\n",
    "\n",
    "joined_table.to_csv('/mnt/tess/astronet/tces-s33_cam2ccd14_sample.csv')\n",
    "\n",
    "# python astronet/preprocess/generate_input_records.py --input_tce_csv_file=/mnt/tess/astronet/tces-s33_cam2ccd14_sample.csv --tess_data_dir=/mnt/tess/lc --output_dir=/mnt/tess/astronet/tfrecords-s33-cam2-ccd14-sample --num_shards=1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RA</th>\n",
       "      <th>Dec</th>\n",
       "      <th>Tmag</th>\n",
       "      <th>Epoc</th>\n",
       "      <th>Period</th>\n",
       "      <th>Duration</th>\n",
       "      <th>Transit_Depth</th>\n",
       "      <th>Sectors</th>\n",
       "      <th>star_rad</th>\n",
       "      <th>star_mass</th>\n",
       "      <th>teff</th>\n",
       "      <th>logg</th>\n",
       "      <th>SN</th>\n",
       "      <th>Qingress</th>\n",
       "      <th>disp_E</th>\n",
       "      <th>disp_J</th>\n",
       "      <th>disp_N</th>\n",
       "      <th>disp_S</th>\n",
       "      <th>disp_B</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>tic_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>93327591</th>\n",
       "      <td>86.116540</td>\n",
       "      <td>-18.100505</td>\n",
       "      <td>12.902</td>\n",
       "      <td>2174.24750</td>\n",
       "      <td>0.10416</td>\n",
       "      <td>0.006167</td>\n",
       "      <td>86167.0</td>\n",
       "      <td>5</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.669</td>\n",
       "      <td>4278.0</td>\n",
       "      <td>4.65</td>\n",
       "      <td>127</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32872129</th>\n",
       "      <td>87.156768</td>\n",
       "      <td>-27.340131</td>\n",
       "      <td>12.590</td>\n",
       "      <td>1468.89711</td>\n",
       "      <td>0.51597</td>\n",
       "      <td>0.061500</td>\n",
       "      <td>2980.0</td>\n",
       "      <td>5</td>\n",
       "      <td>2.21</td>\n",
       "      <td>1.340</td>\n",
       "      <td>6517.0</td>\n",
       "      <td>3.88</td>\n",
       "      <td>11</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37654534</th>\n",
       "      <td>93.266219</td>\n",
       "      <td>-25.821703</td>\n",
       "      <td>12.826</td>\n",
       "      <td>1468.76195</td>\n",
       "      <td>0.21250</td>\n",
       "      <td>0.105333</td>\n",
       "      <td>45280.0</td>\n",
       "      <td>5</td>\n",
       "      <td>1.84</td>\n",
       "      <td>1.480</td>\n",
       "      <td>6845.0</td>\n",
       "      <td>4.08</td>\n",
       "      <td>118</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>299584180</th>\n",
       "      <td>96.837324</td>\n",
       "      <td>-22.523805</td>\n",
       "      <td>12.567</td>\n",
       "      <td>1468.67761</td>\n",
       "      <td>0.48541</td>\n",
       "      <td>0.094000</td>\n",
       "      <td>1480.0</td>\n",
       "      <td>5</td>\n",
       "      <td>1.53</td>\n",
       "      <td>1.380</td>\n",
       "      <td>6615.0</td>\n",
       "      <td>4.21</td>\n",
       "      <td>9</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48621176</th>\n",
       "      <td>96.754353</td>\n",
       "      <td>-24.210266</td>\n",
       "      <td>12.613</td>\n",
       "      <td>2201.74028</td>\n",
       "      <td>0.15972</td>\n",
       "      <td>0.007167</td>\n",
       "      <td>2777.0</td>\n",
       "      <td>5</td>\n",
       "      <td>1.70</td>\n",
       "      <td>1.630</td>\n",
       "      <td>7239.0</td>\n",
       "      <td>4.19</td>\n",
       "      <td>10</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  RA        Dec    Tmag        Epoc   Period  Duration  \\\n",
       "tic_id                                                                   \n",
       "93327591   86.116540 -18.100505  12.902  2174.24750  0.10416  0.006167   \n",
       "32872129   87.156768 -27.340131  12.590  1468.89711  0.51597  0.061500   \n",
       "37654534   93.266219 -25.821703  12.826  1468.76195  0.21250  0.105333   \n",
       "299584180  96.837324 -22.523805  12.567  1468.67761  0.48541  0.094000   \n",
       "48621176   96.754353 -24.210266  12.613  2201.74028  0.15972  0.007167   \n",
       "\n",
       "           Transit_Depth  Sectors  star_rad  star_mass    teff  logg   SN  \\\n",
       "tic_id                                                                      \n",
       "93327591         86167.0        5      0.64      0.669  4278.0  4.65  127   \n",
       "32872129          2980.0        5      2.21      1.340  6517.0  3.88   11   \n",
       "37654534         45280.0        5      1.84      1.480  6845.0  4.08  118   \n",
       "299584180         1480.0        5      1.53      1.380  6615.0  4.21    9   \n",
       "48621176          2777.0        5      1.70      1.630  7239.0  4.19   10   \n",
       "\n",
       "           Qingress  disp_E  disp_J  disp_N  disp_S  disp_B  \n",
       "tic_id                                                       \n",
       "93327591        0.0       0       0       0       0       0  \n",
       "32872129        0.0       0       0       0       0       0  \n",
       "37654534        0.0       0       0       0       0       0  \n",
       "299584180       0.0       0       0       0       0       0  \n",
       "48621176        0.0       0       0       0       0       0  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.set_option('display.max_columns', None)\n",
    "joined_table.sample(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_tces_old():\n",
    "    tceold = pd.read_csv('/mnt/tess/astronet/tces.csv', header=0).set_index('tic_id')\n",
    "\n",
    "    # Only keep the max sectors read.\n",
    "    maxsect = tceold.groupby('tic_id')['Sectors'].max()\n",
    "    tceold = tceold.join(maxsect, on='tic_id', how='right', rsuffix='_max')\n",
    "    tceold = tceold[tceold.Sectors == tceold.Sectors_max]\n",
    "\n",
    "    # Then keep the max row ID.\n",
    "    maxrowid = tceold.groupby('tic_id')['row_id'].max()\n",
    "    tceold = tceold.join(maxrowid, on='tic_id', how='right', rsuffix='_max')\n",
    "    tceold = tceold[tceold.row_id == tceold.row_id_max]\n",
    "\n",
    "    return tceold\n",
    "\n",
    "def generate_tce_bls_instar():\n",
    "    tcenew = pd.read_csv('/mnt/tess/labels/tce_bls_instar.csv', header=0).set_index('tic_id')\n",
    "    tceold = load_tces_old()\n",
    "    tcenorth = pd.read_csv('/mnt/tess/labels/tce_north_instar.csv', header=0).set_index('tic_id')\n",
    "\n",
    "    # Copy from old data where it's missing from the new.\n",
    "    alltce = tcenew.join(tceold, how='outer', on='tic_id', rsuffix='_old')\n",
    "    alltce = alltce.set_index('tic_id')\n",
    "\n",
    "    alltce = alltce.drop(columns=['row_id'])\n",
    "\n",
    "    def fillna(df, col_name):\n",
    "        df.loc[df[col_name].isna(), col_name] = df.loc[df[col_name].isna(), col_name + '_old']\n",
    "\n",
    "    fillna(alltce, 'toi_id')\n",
    "    fillna(alltce, 'Disposition')\n",
    "    fillna(alltce, 'RA')\n",
    "    fillna(alltce, 'Dec')\n",
    "    fillna(alltce, 'Tmag')\n",
    "    fillna(alltce, 'Epoc')\n",
    "    fillna(alltce, 'Period')\n",
    "    fillna(alltce, 'Duration')\n",
    "    fillna(alltce, 'Transit_Depth')\n",
    "    fillna(alltce, 'Sectors')\n",
    "    fillna(alltce, 'camera')\n",
    "    fillna(alltce, 'ccd')\n",
    "    fillna(alltce, 'star_rad')\n",
    "    fillna(alltce, 'star_mass')\n",
    "    fillna(alltce, 'teff')\n",
    "    fillna(alltce, 'logg')\n",
    "    fillna(alltce, 'SN')\n",
    "    fillna(alltce, 'Qingress')\n",
    "\n",
    "    alltce = alltce.drop(columns=[c for c in alltce.columns if c.endswith('_old')])\n",
    "    \n",
    "    alltce = alltce.append(tcenorth)\n",
    "    \n",
    "    alltce['Ilabel'] = alltce['Ilabel'].fillna(False)\n",
    "\n",
    "    alltce.to_csv('/mnt/tess/labels/tce_bls_instar+old.csv')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
